# data_loader.py

import pandas as pd
import os
from datetime import datetime
from glob import glob


class DataLoader:
    """
    负责加载和预处理币安的K线和资金费率数据。

    核心功能:
    - 根据指定的币种、时间范围和周期加载K线数据。
    - 根据指定的币种和时间范围加载资金费率数据。
    - 对数据进行预处理，如计算衍生列、转换时间格式。
    - 将K线和资金费率数据按时间对齐合并，生成一个多重索引的DataFrame。
    """

    def _get_file_paths_for_date_range(self, base_path_template: str, start_date: str, end_date: str) -> list:
        """根据日期范围生成需要读取的月份文件路径列表。"""
        date_range = pd.date_range(start=start_date, end=end_date, freq='MS')
        paths = []
        for dt in date_range:
            path = base_path_template.format(YYYY=dt.year, MM=f"{dt.month:02d}")
            # 使用 glob 应对文件不存在的情况
            found_files = glob(path)
            paths.extend(found_files)
        return sorted(list(set(paths)))

    def load_kline_data(self, symbols: list, start_date: str, end_date: str, timeframe: str) -> pd.DataFrame:
        """
        加载多个币种在指定时间范围内的K线数据。

        :param symbols: 币种列表, e.g., ['BTCUSDT', 'ETHUSDT']
        :param start_date: 开始日期, e.g., '2023-01-01'
        :param end_date: 结束日期, e.g., '2023-06-30'
        :param timeframe: 时间周期, e.g., '1d', '4h'
        :return: 一个包含所有币种K线数据的多重索引DataFrame。
        """
        all_kline_data = []
        kline_cols = [
            'open_time', 'open', 'high', 'low', 'close', 'volume', 'close_time',
            'quote_volume', 'count', 'taker_buy_volume', 'taker_buy_quote_volume', 'ignore'
        ]

        for symbol in symbols:
            path_template = os.path.join(
                'F:\\', 'personal', 'binance_klines', symbol, timeframe,
                f'{symbol}-{timeframe}-{{YYYY}}-{{MM}}.csv'
            )
            print(path_template)
            file_paths = self._get_file_paths_for_date_range(path_template, start_date, end_date)
            print(file_paths)
            if not file_paths:
                print(f"警告: 未找到 {symbol} 在 {start_date} 到 {end_date} 期间的K线数据。")
                continue
            symbol_df = pd.concat([pd.read_csv(f, names=kline_cols, header=0) for f in file_paths], ignore_index=True)

            # --- 数据预处理 ---
            symbol_df['open_time'] = pd.to_datetime(symbol_df['open_time'], unit='ms')
            symbol_df.set_index('open_time', inplace=True)

            # 过滤日期范围
            symbol_df = symbol_df.loc[start_date:end_date]

            # 类型转换
            for col in ['open', 'high', 'low', 'close', 'volume', 'quote_volume', 'taker_buy_volume',
                        'taker_buy_quote_volume']:
                symbol_df[col] = pd.to_numeric(symbol_df[col])

            # 计算衍生数据
            symbol_df['volume_delta'] = 2 * symbol_df['taker_buy_volume'] - symbol_df['volume']

            symbol_df['symbol'] = symbol
            all_kline_data.append(symbol_df)

        if not all_kline_data:
            raise FileNotFoundError("错误: 所有指定币种和时间范围内的数据均未找到。")

        final_df = pd.concat(all_kline_data)
        final_df.set_index(['symbol', final_df.index], inplace=True)
        final_df.rename_axis(['symbol', 'timestamp'], inplace=True)
        final_df.sort_index(inplace=True)
        return final_df

    def load_funding_rate_data(self, symbols: list, start_date: str, end_date: str) -> pd.DataFrame:
        """
        加载多个币种在指定时间范围内的资金费率数据。
        """
        all_funding_data = []

        for symbol in symbols:
            path_template = os.path.join(
                'F:', 'personal', 'binance_data', symbol, 'fundingrate',
                f'{symbol}-fundingRate-{{YYYY}}-{{MM}}.csv'
            )
            file_paths = self._get_file_paths_for_date_range(path_template, start_date, end_date)

            if not file_paths:
                print(f"警告: 未找到 {symbol} 在 {start_date} 到 {end_date} 期间的资金费率数据。")
                continue

            symbol_df = pd.concat([pd.read_csv(f) for f in file_paths], ignore_index=True)

            # --- 数据预处理 ---
            symbol_df.rename(columns={'calc_time': 'timestamp', 'funding_rate': 'fundingRate'}, inplace=True)
            symbol_df['timestamp'] = pd.to_datetime(symbol_df['timestamp'])
            symbol_df.set_index('timestamp', inplace=True)
            symbol_df = symbol_df.loc[start_date:end_date]
            symbol_df['symbol'] = symbol
            all_funding_data.append(symbol_df[['symbol', 'fundingRate']])

        if not all_funding_data:
            return pd.DataFrame()  # 返回空DataFrame

        final_df = pd.concat(all_funding_data)
        final_df.set_index(['symbol', final_df.index], inplace=True)
        final_df.sort_index(inplace=True)
        return final_df

    def load_and_merge_data(self, symbols: list, start_date: str, end_date: str, timeframe: str) -> pd.DataFrame:
        """
        加载、预处理并合并K线和资金费率数据。
        """
        print("开始加载K线数据...")
        kline_data = self.load_kline_data(symbols, start_date, end_date, timeframe)
        print("K线数据加载完成。")

        print("开始加载资金费率数据...")
        funding_data = self.load_funding_rate_data(symbols, start_date, end_date)
        print("资金费率数据加载完成。")

        if funding_data.empty:
            print("警告: 资金费率数据为空，将仅使用K线数据。")
            kline_data['fundingRate'] = 0.0
            return kline_data

        print("开始合并数据...")
        # 使用 merge_asof 进行前向填充合并
        merged_df = pd.merge_asof(
            kline_data.sort_index(level='timestamp'),
            funding_data.sort_index(level='timestamp'),
            on='timestamp',
            by='symbol',
            direction='forward'
        )
        # 如果第一个K线之前没有资金费率，会出现NaN，用0填充
        merged_df['fundingRate'].fillna(0.0, inplace=True)
        print("数据合并完成。")

        return merged_df